import numpy as np
import pandas as pd

from air_web.data_platform import init_db


class GetTradeMap:
    """
    专变:
      第5层作为小行业，第4层作为大行业
      第3层的城镇居民/乡村居民/趸售特殊添加进来(高压居民-城镇居民、高压居民-农村居民、趸售-趸售)
    公变:
      城市台区-城市居民、农村台区-农村居民
    """

    TYPE_MAP_COLS = ["type_id", "type_code", "p_type_id", "type_level"]
    RELA_COLS = [
        "trade_code",
        "trade_id",
        "trade_name",
        "p_trade_code",
        "trade_level",
        "sort_no",
        "type_id",
        "type_code",
        "pare_type_id",
        "pare_type_code",
        "child_trade_code",
    ]

    def __init__(self, table_name, trade_file):
        self.sql_engine = init_db()
        self.table_name = table_name
        self.trade_file = trade_file

        self.trade_df = pd.DataFrame()

    def gen_type_map(self, type_5_df, type_4_df):
        type_5_df = type_5_df.rename(
            columns={
                "trade_id": "type_id",
                "p_trade_id": "p_type_id",
                "trade_name": "type_code",
            }
        )
        type_5_df["type_level"] = 2

        type_4_df["p_type_id"] = 1
        type_4_df["type_level"] = 1
        type_4_df = type_4_df.rename(
            columns={"trade_id": "type_id", "trade_name": "type_code"}
        )

        res_df = pd.concat(
            [type_5_df[self.TYPE_MAP_COLS], type_4_df[self.TYPE_MAP_COLS]]
        )

        other = [
            [20101, "城市居民", 201, 2],
            [20201, "农村居民", 202, 2],
            [201, "城市台区", 2, 1],
            [202, "农村台区", 2, 1],
            [1, "高压用户", 0, 0],
            [2, "低压用户", 0, 0],
        ]
        other_df = pd.DataFrame(other, columns=self.TYPE_MAP_COLS)
        res_df = pd.concat([res_df, other_df])
        self.sql_engine.update_df_by_id(res_df, "type_map")

    def gen_rela(self, type_5_df, type_4_df, trade_4_df):
        rela_df = self.trade_df.loc[
            (self.trade_df["trade_level"] > 4)
            | (self.trade_df["p_trade_code2"].isin(["gaoya", "dunshou"]))
        ]
        # # 第3层的特殊加入
        # spec_df = self.trade_df.loc[self.trade_df['p_trade_code2'] == '9900']
        # # spec_df['level_5_code'] = spec_df['trade_code']
        # # spec_df['level_4_code'] = spec_df['p_trade_code2']
        # rela_df = pd.concat([spec_df, rela_df])

        level_5_df = type_5_df[
            ["trade_code", "trade_id", "trade_name"]
        ].rename(
            columns={
                "trade_code": "level_5_code",
                "trade_id": "type_id",
                "trade_name": "type_code",
            }
        )
        rela_df = rela_df.merge(level_5_df, on="level_5_code", how="left")

        level_4_df = type_4_df[
            ["trade_code", "trade_id", "trade_name"]
        ].rename(
            columns={
                "trade_code": "level_4_code",
                "trade_id": "pare_type_id",
                "trade_name": "pare_type_code",
            }
        )
        rela_df = rela_df.merge(level_4_df, on="level_4_code", how="left")

        type_rela_df = rela_df.loc[rela_df["trade_level"].isin([3, 5])]
        type_rela_df = type_rela_df.merge(
            type_5_df[["trade_code", "trade_id"]], on="trade_code", how="left"
        )

        new_other_rela_df = pd.DataFrame()
        other_rela_df = rela_df.loc[~rela_df["trade_level"].isin([3, 4, 5])]
        for type_id, group_df in other_rela_df.groupby("type_id"):
            group_df["trade_id"] = (
                group_df["trade_code"].factorize()[0] + type_id * 1000 + 1
            )
            new_other_rela_df = pd.concat([new_other_rela_df, group_df])

        res_df = pd.concat(
            [type_rela_df[self.RELA_COLS], new_other_rela_df[self.RELA_COLS]]
        )

        # 给帆软做4，5，6层的树形下拉框使用，添加第4层
        trade_4_df = trade_4_df[self.RELA_COLS[0:5]]
        for col in self.RELA_COLS[5:]:
            trade_4_df[col] = np.nan
        res_df = pd.concat([trade_4_df, res_df])

        res_df = res_df.replace({np.nan: None})
        self.sql_engine.update_df_by_id(res_df, self.table_name)

    def add_level_6_code(self):
        six_seven_dict = (
            self.trade_df.loc[
                self.trade_df["trade_level"] == 7,
                ["trade_code", "p_trade_code"],
            ]
            .set_index("trade_code")["p_trade_code"]
            .to_dict()
        )

        self.trade_df.loc[
            self.trade_df["trade_level"] == 6, "child_trade_code"
        ] = self.trade_df["trade_code"]
        self.trade_df.loc[
            self.trade_df["trade_level"] == 7, "child_trade_code"
        ] = self.trade_df["trade_code"].replace(six_seven_dict)
        self.trade_df.loc[
            self.trade_df["trade_level"] == 8, "child_trade_code"
        ] = self.trade_df["p_trade_code"].replace(six_seven_dict)

    def get_trade_df(self):
        self.trade_df = pd.read_csv(self.trade_file, header=0)

    def main(self):
        self.get_trade_df()
        self.add_level_6_code()

        self.trade_df["p_trade_code2"] = self.trade_df["p_trade_code"]

        self.trade_df.loc[
            self.trade_df["trade_code"] == "9999", "p_trade_code2"
        ] = "dunshou"
        self.trade_df.loc[
            self.trade_df["trade_code"].isin(["9910", "9920"]), "p_trade_code2"
        ] = "gaoya"
        self.trade_df.loc[
            self.trade_df["p_trade_code2"].isin(["gaoya", "dunshou"]),
            "level_5_code",
        ] = self.trade_df["trade_code"]
        self.trade_df.loc[
            self.trade_df["p_trade_code2"].isin(["gaoya", "dunshou"]),
            "level_4_code",
        ] = self.trade_df["p_trade_code2"]

        # 第5层
        type_5_df = self.trade_df.loc[self.trade_df["trade_level"] == 5]
        # 第3层的特殊加入
        type_3_df = self.trade_df.loc[
            self.trade_df["p_trade_code2"].isin(["gaoya", "dunshou"])
        ]
        # type_3_df['level_5_code'] = type_3_df['trade_code']
        # type_3_df['level_4_code'] = type_3_df['p_trade_code2']
        type_5_df = pd.concat([type_3_df, type_5_df])
        # 给第5层（小行业）编码和父级4层（大行业）编码
        type_5_df["p_trade_id"] = (
            type_5_df["p_trade_code2"].factorize()[0] + 101
        )
        new_type_5_df = pd.DataFrame()
        for para_type_id, group_df in type_5_df.groupby("p_trade_id"):
            group_df["trade_id"] = (
                group_df["trade_code"].factorize()[0] + para_type_id * 100 + 1
            )
            new_type_5_df = pd.concat([new_type_5_df, group_df])

        # 第4层
        type_4_map = (
            new_type_5_df[["p_trade_code2", "p_trade_id"]]
            .set_index("p_trade_code2")["p_trade_id"]
            .to_dict()
        )
        type_4_df = self.trade_df.loc[self.trade_df["trade_level"] == 4]
        # 第3层的特殊加入
        type_3_df["trade_code"] = type_3_df["p_trade_code2"]
        type_3_df.loc[
            type_3_df["trade_code"] == "gaoya", "trade_name"
        ] = "高压居民"
        type_4_df = pd.concat([type_3_df, type_4_df])
        # 替换第4层的编码
        type_4_df["trade_id"] = type_4_df["trade_code"].replace(type_4_map)

        # 给帆软做4，5，6层的树形下拉框使用，添加第4层
        trade_4_df = self.trade_df.loc[self.trade_df["trade_level"] == 4]
        trade_4_df["trade_id"] = type_4_df["trade_code"].replace(type_4_map)

        # self.gen_type_map(new_type_5_df, type_4_df)
        self.gen_rela(new_type_5_df, type_4_df, trade_4_df)


if __name__ == "__main__":
    table_name = "trade_code_id_rela"
    trade_file = "/home/smxu/hebei/chongqing_trade_code.csv"

    gtm = GetTradeMap(table_name, trade_file)
    gtm.main()
